import argparse
import json
import os


from driver.Dataloader import read_corpus
from collections import Counter


relative_distances = [ str(dis) for dis in range(1, 16)]

relations = [
    'Question-answer_pair', 'Comment', 'Acknowledgement',
    'Continuation', 'Elaboration', 'Q-Elab', 
    'Contrast', 'Explanation', 'Clarification_question', 
    'Result', 'Correction', 'Parallel', 
    'Conditional', 'Alternation', 'Narration', 'Background'
]

def max_r(file_path):
    with open(file_path, mode='r', encoding='utf8') as inf:
        data = json.load(inf)

    r_dis = Counter()
    for inst in  data:
        for rel in inst['relations']:
            r_dis[ int(rel['y']) - int(rel['x'])] += 1
    r_dis.most_common()            
    print()

def load_examples(file_path):
    examples = []
    with open(file_path, encoding='utf8', mode='r') as inf:
        data = json.load(inf)

        print("instance num: ", len(data))
        input_str = ""

        for inst in data:

            for rel in inst['relations']:
                cur_idx = rel['y']
                linked_idx = rel['x']
                if cur_idx > linked_idx:
                    
                    edus = ["<u-1> root: root"]
                    for idx in range(cur_idx + 1):
                        edu = inst['edus'][idx]
                        edu_text  = "<u" +  str(idx) + "> "  +  edu['speaker']  + ": " + edu['text']
                        edus.append(edu_text)
                    
                    input_str = " ".join(edus)
                    
                    r_distance = linked_idx
                    sdrt_relation = rel['type']

                    output_str = str(r_distance) + " " + sdrt_relation

                    examples.append( dict( instruction=instruction_str, input=input_str, output=output_str) ) 
    return examples

if __name__ == '__main__':
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--sdrt_dir', default='data/STAC_sp')
    argparser.add_argument('--out', default='STAC_train')
    args, extra_args = argparser.parse_known_args()

    instruction_str = \
        "Please perform SDRT parsing to the index of linked utterance " + \
        "and select the relation between them from the following relation set <" +  " ".join(relations) + ">"
    
    train_path = os.path.join(args.sdrt_dir, "train.json")
    dev_path = os.path.join(args.sdrt_dir, "dev.json")

    examples = load_examples(train_path) + load_examples(dev_path)

    outfile = args.out + '.json'
    with open(outfile, 'w') as f:
        json.dump(examples, f, indent=4, ensure_ascii=False)
    print("convert to: ", outfile)
    print("num: ", len(examples))